#!/bin/sh

# MTD partitions
#Device       Name      Start    Size(B)  Size(MiB)  Size(Blks)
#mtd0        spl        0x0   0x200000          2          16
#mtd1      uboot   0x200000   0x200000          2          16
#mtd2  uboot env   0x400000   0x200000          2          16
#mtd3      linux   0x600000   0xc00000         12          96
#mtd4     rootfs  0x1200000  0x2d00000         45         360
#mtd5       data  0x3f00000  0x4100000         65         520
#total size = 0x8000000, 128MiB

UPDATER_VERSION="20"

source /platform

speak() {
  if [ "$UPDATER_SPEECH" = "1" ]; then
    msg "speak $1"
    espeak "$1"
  fi
}

msg() {
  echo "$@" >/dev/console
}

# Prints information
msg_splash() {
  msg $1

  if [ "$splash_enabled" = 1 ]; then
    progress=$2

    if [ "$progress" = "" ]; then
      progress="0"
    fi

    TMPDIR=$SPLASH_DIR psplash-write "MSG $1"
    TMPDIR=$SPLASH_DIR psplash-write "PROGRESS $progress"
    usleep 500000
  fi
}

SPLASH_DIR="/mnt/.splash"

start_splash() {
  if [ -e /usr/bin/psplash ]; then 
    splash_enabled=1
    # set up psplash
    mkdir -p "$SPLASH_DIR"
    mount tmpfs -t tmpfs $SPLASH_DIR -o,size=40k 2>&1 > /dev/console
    TMPDIR=$SPLASH_DIR psplash& 2>&1 > /dev/console
    sleep 1
  fi
}

BOOT="mtd1"
KERNEL="mtd3"
ROOTFS="mtd4"
DATAPART="mtd5"
ROOTFS_MOUNT="/root"

MTD_ROOTFS=4
MTD_DATA=5

format_data_partition() {
  msg_splash "Formatting data partition"
  ubidetach -m $MTD_DATA 2>/dev/null
  flash_eraseall /dev/mtd$MTD_DATA
  if ! ubiattach -m $MTD_DATA; then
    msg_splash "Error attaching to MTD data"
    return 1
  fi

  ubimkvol /dev/ubi1 -N ubidata -m || return 1
}

mount_data() {
  DATA_MOUNT_POINT=/data

  msg "mounting data partition"
  mkdir -p $DATA_MOUNT_POINT
  if ! mount -t ubifs ubi1:ubidata $DATA_MOUNT_POINT; then
    msg "data partition mount failed, formatting"
    format_data_partition
    if ! mount -t ubifs ubi1:ubidata $DATA_MOUNT_POINT; then
      msg_splash "data partition mount after format failed"
      speak "updater, data partition error, please service unit"
      return 1
    fi
  fi

  mkdir -p $DATA_MOUNT_POINT/log

  msg "data partition mounted"

  return 0
}

update_uboot() {
  if [ -e u-boot.img ]; then
    msg "checking if u-boot needs updated"
    SIZE=$(stat -c%s u-boot.img)
    nanddump -l $SIZE -f uboot.dump /dev/$BOOT || return 1
    dd if=uboot.dump of=uboot.dump2 bs=$SIZE count=1
    if ! diff u-boot.img uboot.dump2 2>/dev/null; then
      msg_splash "updating uboot"
      flash_eraseall /dev/$BOOT
      nandwrite -p /dev/$BOOT u-boot.img
      boot_updated=1
    else
      msg "u-boot is up to date"
    fi
    rm uboot.dump
    rm uboot.dump2
  fi
}

update_kernel() {
  kernel_needs_update=0

  if [ -e zImage ]; then
    msg "checking if kernel needs to be updated"
    SIZE=$(stat -c%s zImage)
    nanddump -l $SIZE -f zImage.dump /dev/$KERNEL || return 1
    dd if=zImage.dump of=zImage.dump2 bs=$SIZE count=1
    if ! diff zImage zImage.dump2 2>/dev/null; then
      msg "zImage needs updated"
      kernel_needs_update=1
    else
      msg "kernel is up to date"
    fi
    rm zImage.dump
    rm zImage.dump2
  fi

  if dmesg | grep "Machine.*ULL"; then
    DTB_FILE=zImage-ull.dtb
  else
    DTB_FILE=zImage.dtb
  fi

  if [ -e $DTB_FILE ]; then
    msg "checking if dtb needs to be updated"
    SIZE=$(stat -c%s $DTB_FILE)
    nanddump -s 0xbe0000 -l $SIZE -f dtb.dump /dev/$KERNEL || return 1
    dd if=dtb.dump of=dtb.dump2 bs=$SIZE count=1
    if ! diff $DTB_FILE dtb.dump2 2>/dev/null; then
      msg "dtb needs updated"
      kernel_needs_update=1
    else
      msg "dtb is up to date"
    fi
    rm dtb.dump
    rm dtb.dump2
  fi

  if [ "$kernel_needs_update" = "1" ]; then
    msg_splash "updating kernel image"
    msg "DTB: $DTB_FILE"
    flash_eraseall /dev/$KERNEL || return 1
    msg "writing zImage"
    nandwrite -p /dev/$KERNEL zImage || return 1
    msg "writing dtb"
    nandwrite -p -s 0xbe0000 /dev/$KERNEL $DTB_FILE || return 1
    msg_splash "kernel update complete"
    boot_updated=1
  fi

  return 0
}

update_rootfs() {
  if [ -e rootfs.ubi ]; then
    msg_splash "Writing new image, please wait ..."
    ubidetach -m $MTD_ROOTFS 2>/dev/null
    flash_eraseall /dev/$ROOTFS || return 1
    ubiformat /dev/$ROOTFS -f rootfs.ubi -s 2048 -O 2048
    msg_splash "Image update complete"
  fi
  return 0
}

process_update() {
  UPDATE_FILE=$1
  if ! cpio -id < $UPDATE_FILE; then
    msg_splash "Failed to extract update image"
    return 1
  fi

  cd /tmp
  if ! md5sum -c update.md5; then
    msg_splash "Error: checksum error in update file"
    return 1
  else
    speak "updating system"
    boot_updated=0
    update_uboot || return 1
    update_kernel || return 1

    if [ "$boot_updated" = "1" ]; then
      msg_splash "Boot components updated, rebooting ..."
      speak "rebooting"
      reboot -f
    fi

    update_rootfs || return 1

    cp version.txt /data/
  fi
  cd /
  # wait for user to remove USB flash disk
  # not doing this for now
  # while [ -e /sys/class/scsi_disk/* ]; do sleep 1; done

  speak "update complete"
  return 0
}

initialize() {
  msg "============================================================"
  msg "Updater version $UPDATER_VERSION"

  mount -t devtmpfs none /dev
  mount -t sysfs sysfs /sys
  mount -t proc proc /proc

  start_splash 

  plat_init

  ubidetach -m $MTD_ROOTFS
  ubidetach -m $MTD_DATA

  if ! ubiattach -m $MTD_ROOTFS; then
    flash_eraseall /dev/mtd$MTD_ROOTFS
    ubiattach -m $MTD_ROOTFS
  fi

  if ! ubiattach -m $MTD_DATA; then
    flash_eraseall /dev/mtd$MTD_DATA
    ubiattach -m $MTD_DATA
  fi

  msg "Sleeping for 3 second(s) for USB flash to settle..."

  speak "updater version $UPDATER_VERSION"

  sleep 3

  mkdir -p /usb
  mkdir -p /data
}

mount_usb() {
  if mount /dev/sda1 /usb 2>/dev/null; then
    usb_mounted=1
    msg 'found disk at /dev/sda1'
    return 0
  fi

  if [ "$usb_mounted" = "0" ]; then
    if mount /dev/sda /usb 2>/dev/null; then
      usb_mounted=1
      msg 'found disk at /dev/sda'
      return 0
    fi
  fi

  return 1
}

mount_sd_boot() {
  if mount /dev/mmcblk0p1 /boot 2>/dev/null; then
    msg 'found sd boot partition'
    return 0
  fi

  return 1
}

umount_usb() {
  if mount | grep usb; then
    umount /usb
  fi

  return 0
}

# searches for update*.pwu and update*.img files
# and then uses the newest version file it finds
# also supports the legacy update.img format for now

find_update_file() {
  DIR=$1
  cd $DIR
  extensions="pwu img"
  update=""
  for ext in $extensions; do
    update=$(echo $(ls update_dart*.${ext} -v -r 2>/dev/null) | sed "s/update.${ext}//" | sed "s/${ext}.*/${ext}/" | sed "s/^\s*//")
    if echo $update | grep update >/dev/null; then
      echo $update
      break
    fi
  done

  if [ "${update}" = "" ]; then
    if [ -e update.img ]; then
      echo update.img
    fi
  fi
  cd - >/dev/null
}

update_from_usb() {
  speak "updater, found u s b disk"
  update_file=$(find_update_file /usb)
  if echo $update_file | grep update; then
    msg_splash "USB update: $update_file"
    if ! process_update /usb/$update_file; then
      msg_splash "Failed to process update from USB"
      speak "updater, failed to process update from u s b"
      return 1
    else
      msg_splash "Update from USB complete"
      return 0
    fi
  fi

  return 1
}

update_from_sd() {
  update_file=$(find_update_file /boot)
  if echo $update_file | grep update; then
    msg_splash "SD update: $update_file"
    if ! process_update /boot/$update_file; then
      msg_splash "Failed to process update from SD"
      speak "updater, failed to process update from s d"
      return 1
    else
      msg_splash "Update complete, please remove card and cycle power"
      speak "update from s d complete, please remove s d card and cycle power"
      sleep 9999d
    fi
  fi
  return 1
}

update_from_data() {
  update_file=$(find_update_file /data/update)
  if echo $update_file | grep update; then
    msg_splash "Data partition update: $update_file"
    if ! process_update /data/update/$update_file; then
      msg_splash "Failed to process update from nand"
      speak "updater, failed to process update from nand"
    else
      msg_splash "Update from nand complete"
    fi

    msg_splash "Removing update files from data partition"
    rm /data/update/*
  fi
}

#unmount_usb() {
#}

mount_rootfs_sd() {
  if ! mount /dev/mmcblk0p2 $ROOTFS_MOUNT; then
    msg_splash "Error mount SD rootfs, please fix ..."
    speak "updater, S D file system error"
    sleep 9999d
  fi

  mkdir -p $ROOTFS_MOUNT/media/boot
  mount --move /boot $ROOTFS_MOUNT/media/boot
}

mount_rootfs_nand() {
  # may be attached by kernel on boot, but we have
  # this here in case we programmed the rootfs
  ubiattach -m $MTD_ROOTFS

  if ! mount -tubifs ubi0:rootfs $ROOTFS_MOUNT; then
    msg_splash "Error mounting nand rootfs, please fix ..."
    speak "updater, root f s file system error, please service unit"
    sleep 9999d
  fi
}

boot() {
  mkdir -p $ROOTFS_MOUNT

  if [ -e /dev/mmcblk0p1 ]; then
    msg_splash "Booting from SD ..."
    speak "booting system from S D"
    mount_rootfs_sd
  else
    msg_splash "Booting from NAND ..."
    speak "booting system"
    mount_rootfs_nand
  fi

  umount_usb

  msg "Moving filesystems into rootfs..."
  mkdir -p $ROOTFS_MOUNT/media/data
  mount --move /data $ROOTFS_MOUNT/media/data
  mount --move /dev $ROOTFS_MOUNT/dev
  mount --move /proc $ROOTFS_MOUNT/proc
  mount --move /sys $ROOTFS_MOUNT/sys

  msg_splash "switching to main filesystem"

  exec switch_root -c /dev/console $ROOTFS_MOUNT /sbin/init 5
}

update() {
  # update precedence usb, SD, data partition
  msg "checking usb ..."
  if ! (mount_usb && update_from_usb); then
    msg "checking sd ..."
    if ! (mount_sd_boot && update_from_sd); then
      msg "checking data ..."
      update_from_data
    fi
  fi
}

initialize
mount_data
update
boot
